package weiyao.Detector;

import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.selector.Html;
import weiyao.GUI.ScanResultCallback;
import weiyao.Vul.VulnerabilityConfig;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.net.URLEncoder;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class SqlInjectionDetector extends BaseDetector {

    // 常见URL参数名（用于智能探测）
    private static final String[] COMMON_URL_PARAMS = {
            "id", "user", "username", "name", "email", "product",
            "category", "page", "view", "search", "q", "query",
            "sort", "order", "limit", "offset", "key", "value"
    };

    // 提取登录信息的正则表达式（优化版）
    private static final Pattern LOGIN_PATTERN =
            Pattern.compile("Your Login name:(.*?)<br>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL);

    // 提取密码信息的正则表达式（优化版）
    private static final Pattern PASSWORD_PATTERN =
            Pattern.compile("Your Password:(.*?)</font>", Pattern.CASE_INSENSITIVE | Pattern.DOTALL);

    @Override
    public void detect(Page page, String baseUrl, ScanResultCallback callback) {
        // 1. 检测URL参数注入
        detectUrlParamInjection(page, baseUrl, callback);

        // 2. 检测表单参数注入
        detectFormInjection(page, baseUrl, callback);
    }

    /**
     * 检测URL参数中的SQL注入漏洞
     */
    private void detectUrlParamInjection(Page page, String baseUrl, ScanResultCallback callback) {
        String currentUrl = page.getUrl().get();

        // 提取现有URL参数
        Map<String, String> existingParams = extractUrlParams(currentUrl);

        // 1. 测试现有参数
        testExistingUrlParams(currentUrl, existingParams, callback);

        // 2. 智能探测常见参数
        probeCommonUrlParams(currentUrl, callback);
    }

    /**
     * 检测表单参数中的SQL注入漏洞
     */
    private void detectFormInjection(Page page, String baseUrl, ScanResultCallback callback) {
        Html html = page.getHtml();
        List<String> formElements = html.xpath("//form").all();

        for (String formHtml : formElements) {
            Html form = new Html(formHtml);
            String action = form.xpath("//form/@action").get() != null ?
                    form.xpath("//form/@action").get() : baseUrl;
            String method = form.xpath("//form/@method").get() != null ?
                    form.xpath("//form/@method").get() : "GET";

            String formUrl = makeAbsoluteUrl(baseUrl, action);
            Map<String, String> formParams = collectFormParams(form);

            // 使用重载的 testFormParameters 方法（6个参数）
            testFormParameters(formUrl, method, formParams, VulnerabilityConfig.SQL_PAYLOADS,
                    (url, paramName, payload, response) -> {
                        // 修改点1: 检测响应中是否包含"wtq"
                        if (response.contains("wtq")) {
                            String fullUrl = buildUrlWithParams(url, Collections.singletonMap(paramName, payload));
                            String resp = response.toString();

                            // 提取响应中的login和password信息
                            String loginInfo = extractLoginInfo(resp);
                            String passwordInfo = extractPasswordInfo(resp);

                            // 构建提取信息字符串
                            StringBuilder extractedInfo = new StringBuilder();
                            if (!loginInfo.isEmpty()) {
                                extractedInfo.append("\n提取的登录信息: ").append(loginInfo);
                            }
                            if (!passwordInfo.isEmpty()) {
                                extractedInfo.append("\n提取的密码信息: ").append(passwordInfo);
                            }

                            String details = "表单SQL注入漏洞!\nURL: " + fullUrl +
                                    "\n参数: " + paramName +
                                    "\n载荷: " + payload +
                                    extractedInfo;  // 添加提取的信息

                            callback.onVulnerabilityFound("SQL注入", details);
                            return true;
                        }
                        return false;
                    }, callback);
        }
    }

    // 提取登录信息
    private String extractLoginInfo(String htmlResponse) {
        Matcher matcher = LOGIN_PATTERN.matcher(htmlResponse);
        if (matcher.find()) {
            return matcher.group(1).trim();
        }
        return "";
    }

    // 提取密码信息
    private String extractPasswordInfo(String htmlResponse) {
        Matcher matcher = PASSWORD_PATTERN.matcher(htmlResponse);
        if (matcher.find()) {
            return matcher.group(1).trim();
        }
        return "";
    }

    /**
     * 提取URL中的参数
     */
    private Map<String, String> extractUrlParams(String url) {
        Map<String, String> params = new HashMap<>();
        Pattern pattern = Pattern.compile("([^?=&]+)=([^&]*)");
        Matcher matcher = pattern.matcher(url);

        while (matcher.find()) {
            params.put(matcher.group(1), matcher.group(2));
        }
        return params;
    }

    /**
     * 测试现有URL参数
     */
    private void testExistingUrlParams(String url, Map<String, String> params,
                                       ScanResultCallback callback) {
        if (params.isEmpty()) return;

        // 构建基础URL（不带参数）
        String baseUrl = url.split("\\?")[0];

        // 对每个参数进行注入测试
        for (Map.Entry<String, String> entry : params.entrySet()) {
            String paramName = entry.getKey();
            String originalValue = entry.getValue();

            // 测试每个SQL注入payload
            for (String payload : VulnerabilityConfig.SQL_PAYLOADS) {
                // 创建新参数集（只修改当前参数）
                Map<String, String> testParams = new HashMap<>(params);
                testParams.put(paramName, payload);

                // 构建测试URL
                String testUrl = buildUrlWithParams(baseUrl, testParams);

                // 发送请求并检测响应
                String response = sendGetRequest(testUrl);

                // 修改点2: 检测响应中是否包含"wtq"
                if (response.contains("wtq")) {
                    // 提取响应中的login和password信息
                    String loginInfo = extractLoginInfo(response);
                    String passwordInfo = extractPasswordInfo(response);

                    // 构建提取信息字符串
                    StringBuilder extractedInfo = new StringBuilder();
                    if (!loginInfo.isEmpty()) {
                        extractedInfo.append("\n提取的登录信息: ").append(loginInfo);
                    }
                    if (!passwordInfo.isEmpty()) {
                        extractedInfo.append("\n提取的密码信息: ").append(passwordInfo);
                    }

                    String details = "URL参数SQL注入漏洞!\nURL: " + testUrl +
                            "\n参数: " + paramName +
                            "\n载荷: " + payload +
                            extractedInfo;  // 添加提取的信息

                    callback.onVulnerabilityFound("SQL注入", details);
                }
            }
        }
    }

    /**
     * 智能探测常见URL参数
     */
    private void probeCommonUrlParams(String baseUrl, ScanResultCallback callback) {
        // 构建基础URL（不带参数）
        String cleanUrl = baseUrl.split("\\?")[0];

        // 测试每个常见参数
        for (String paramName : COMMON_URL_PARAMS) {
            // 使用基本探测值（如1）
            Map<String, String> testParams = new HashMap<>();
            testParams.put(paramName, "1");

            // 构建测试URL
            String testUrl = buildUrlWithParams(cleanUrl, testParams);

            // 发送初始请求（检查参数是否被接受）
            String initialResponse = sendGetRequest(testUrl);
            if (!isErrorPage(initialResponse)) {
                // 参数有效，进行注入测试
                for (String payload : VulnerabilityConfig.SQL_PAYLOADS) {
                    testParams.put(paramName, payload);
                    String injectionUrl = buildUrlWithParams(cleanUrl, testParams);
                    String response = sendGetRequest(injectionUrl);

                    // 修改点3: 检测响应中是否包含"wtq"
                    if (response.contains("wtq")) {
                        // 提取响应中的login和password信息
                        String loginInfo = extractLoginInfo(response);
                        String passwordInfo = extractPasswordInfo(response);

                        // 构建提取信息字符串
                        StringBuilder extractedInfo = new StringBuilder();
                        if (!loginInfo.isEmpty()) {
                            extractedInfo.append("\n提取的登录信息: ").append(loginInfo);
                        }
                        if (!passwordInfo.isEmpty()) {
                            extractedInfo.append("\n提取的密码信息: ").append(passwordInfo);
                        }

                        String details = "探测到URL参数SQL注入漏洞!\nURL: " + injectionUrl +
                                "\n参数: " + paramName +
                                "\n载荷: " + payload +
                                extractedInfo;  // 添加提取的信息

                        callback.onVulnerabilityFound("SQL注入", details);
                    }
                }
            }
        }
    }

    /**
     * 构建带参数的URL（带URL编码）
     */
    private String buildUrlWithParams(String baseUrl, Map<String, String> params) {
        StringBuilder urlBuilder = new StringBuilder(baseUrl);
        if (!params.isEmpty()) {
            urlBuilder.append("?");
            boolean first = true;
            for (Map.Entry<String, String> entry : params.entrySet()) {
                if (!first) {
                    urlBuilder.append("&");
                }
                try {
                    urlBuilder.append(URLEncoder.encode(entry.getKey(), "UTF-8"))
                            .append("=")
                            .append(URLEncoder.encode(entry.getValue(), "UTF-8"));
                } catch (Exception e) {
                    // 回退到简单编码
                    urlBuilder.append(entry.getKey())
                            .append("=")
                            .append(entry.getValue());
                }
                first = false;
            }
        }
        return urlBuilder.toString();
    }

    /**
     * 发送GET请求并返回响应内容
     */
    protected String sendGetRequest(String urlString) {
        StringBuilder response = new StringBuilder();
        HttpURLConnection connection = null;

        try {
            URL url = new URL(urlString);
            connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("GET");
            connection.setConnectTimeout(5000); // 5秒连接超时
            connection.setReadTimeout(10000);  // 10秒读取超时

            // 读取响应
            try (BufferedReader reader = new BufferedReader(
                    new InputStreamReader(connection.getInputStream()))) {
                String line;
                while ((line = reader.readLine()) != null) {
                    response.append(line);
                }
            }
        } catch (Exception e) {
            logError("GET请求失败: " + urlString + " - " + e.getMessage());
        } finally {
            if (connection != null) {
                connection.disconnect();
            }
        }

        return response.toString();
    }



    /**
     * 检查是否为错误页面
     */
    private boolean isErrorPage(String response) {
        if (response == null || response.isEmpty()) return true;
        String responseLower = response.toLowerCase();
        return responseLower.contains("404 not found") ||
                responseLower.contains("500 internal server error") ||
                responseLower.contains("error") ||
                responseLower.contains("exception") ||
                responseLower.contains("not found") ||
                responseLower.contains("forbidden");
    }

    /**
     * 收集表单参数
     */
    public Map<String, String> collectFormParams(Html form) {
        Map<String, String> params = new HashMap<>();
        List<String> inputs = form.xpath("//input").all();

        for (String input : inputs) {
            Html inputHtml = new Html(input);
            String name = inputHtml.xpath("//input/@name").get();
            String value = inputHtml.xpath("//input/@value").get();
            String type = inputHtml.xpath("//input/@type").get();

            if (name != null && !"submit".equalsIgnoreCase(type)) {
                params.put(name, value != null ? value : "");
            }
        }
        return params;
    }

    /**
     * 发送POST请求
     */
    private String sendPostRequest(String urlString, Map<String, String> params) {
        StringBuilder response = new StringBuilder();
        HttpURLConnection connection = null;

        try {
            URL url = new URL(urlString);
            connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("POST");
            connection.setDoOutput(true);
            connection.setConnectTimeout(5000);
            connection.setReadTimeout(10000);

            // 构建请求体
            StringBuilder postData = new StringBuilder();
            for (Map.Entry<String, String> param : params.entrySet()) {
                if (postData.length() != 0) postData.append('&');
                postData.append(URLEncoder.encode(param.getKey(), "UTF-8"));
                postData.append('=');
                postData.append(URLEncoder.encode(param.getValue(), "UTF-8"));
            }

            // 发送请求
            try (OutputStream os = connection.getOutputStream()) {
                byte[] input = postData.toString().getBytes("utf-8");
                os.write(input, 0, input.length);
            }

            // 读取响应
            try (BufferedReader br = new BufferedReader(
                    new InputStreamReader(connection.getInputStream(), "utf-8"))) {
                String responseLine;
                while ((responseLine = br.readLine()) != null) {
                    response.append(responseLine);
                }
            }
        } catch (Exception e) {
            logError("POST请求失败: " + urlString + " - " + e.getMessage());
        } finally {
            if (connection != null) {
                connection.disconnect();
            }
        }

        return response.toString();
    }

    /**
     * 记录错误日志
     */
    private void logError(String message) {
        System.err.println("[ERROR] " + message);
    }
}